// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cantor/proto/parser/ast.pb.h>
#include <turbo/container/flat_hash_map.h>
#include <turbo/utility/status.h>
#include <turbo/log/logging.h>
#include <memory>
#include <variant>

#include <cantor/unresolved/expr_traits.h>

#include <cantor/unresolved/functions/func_expr.h>


namespace cantor::unresolved {
    struct DummyExpression {
    };

    using ExprVisitor = std::variant<
        UnresolvedConstExpr,
        UnresolvedRangeVar,
        UnresolvedRowExpr,
        UnresolvedFuncExpr,
        UnresolvedCtorMapFuncExpr,
        UnresolvedCtorStructFuncExpr,
        UnresolvedCtorUnionFuncExpr,
        UnresolvedCtorArrayFuncExpr,
        UnresolvedCtorSetFuncExpr,
        UnresolvedCtorTupleFuncExpr,
        UnresolvedCtorJsonFuncExpr,
        UnresolvedCtorJsonArrayFuncExpr,
        UnresolvedCtorEnumFuncExpr,
        UnresolvedCtorProtoFuncExpr,
        UnresolvedCtorUdfFuncExpr
    >;

    class UnresolvedDispatchExpr {
    public:
        static constexpr cantor::pb::ast::ExprNodeType TYPE = cantor::pb::ast::ExprNodeType::ENT_NONE;

    public:
        UnresolvedDispatchExpr() = default;

        virtual ~UnresolvedDispatchExpr() = default;

        explicit UnresolvedDispatchExpr(const cantor::pb::ast::DispatchExpr *st) : expr(st) {
        }


        void attach(const cantor::pb::ast::DispatchExpr *st) {
            expr = st;
        }

        const cantor::pb::ast::DispatchExpr *expr{nullptr};

    public:
        /// @brief Cast the underlying statement to a concrete typed visitor.
        ///
        /// This verifies that the `stmt->type()` matches `Visitor::TYPE`, and then
        /// constructs the visitor.
        ///
        /// @tparam Visitor A class derived from `UnresolvedVisitor<Visitor>`
        /// @return `turbo::Result<Visitor>` if successful, otherwise contains error.
        template<class Visitor>
        turbo::Result<Visitor> cast_to() const {
            static_assert(IsDispatchExpr<Visitor>::value, "not a dispatch expression");
            if (!expr) {
                return turbo::data_loss_error("no stmt setting");
            }
            if (expr->type() != Visitor::TYPE) {
                return turbo::invalid_argument_error("Failed to cast statement to type - statement type mismatch [",
                                                     cantor::pb::ast::ExprNodeType_Name(expr->type()), ":",
                                                     cantor::pb::ast::ExprNodeType_Name(Visitor::TYPE), "]");
            }
            return Visitor(*expr);
        }

        /// @brief Check if the stored statement matches the `Visitor::TYPE`.
        template<class Visitor>
        bool is_type() const {
            return expr && (expr->type() == Visitor::TYPE || Visitor::TYPE == cantor::pb::ast::ExprNodeType::ENT_NONE);
        }

        /// @brief Check if the stored statement matches the specified statement type.
        template<cantor::pb::ast::ExprNodeType Stype>
        bool is_expr() const {
            return expr && (expr->type() == Stype || Stype == cantor::pb::ast::ExprNodeType::ENT_NONE);
        }

        /// @brief Cast the statement to a visitor or crash fatally on error.
        ///
        /// Use this only when you're confident about the statement type.
        /// Crashes the program with a fatal log on mismatch or null.
        template<class Visitor>
        Visitor cast_or_die() const {
            auto r = cast_to<Visitor>();
            if (!r.ok()) {
                KLOG(FATAL) << r.status();
            }
            return std::move(r).value_or_die();
        }

        turbo::Result<ExprVisitor> cast_to(cantor::pb::ast::ExprNodeType Stype) const {
            switch (Stype) {
                case cantor::pb::ast::ExprNodeType::ENT_CONSTANT_EXPR:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_CONSTANT_EXPR>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_RANGE_VAR:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_RANGE_VAR>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_ROW_EXPR:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_ROW_EXPR>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_FUNC_EXPR:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_FUNC_EXPR>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_MAP_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_MAP_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_STRUCT_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_STRUCT_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_UNION_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_UNION_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_ARRAY_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_ARRAY_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_SET_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_SET_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_TUPLE_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_TUPLE_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_JSON_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_JSON_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_JSON_ARRAY_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_JSON_ARRAY_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_ENUM_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_ENUM_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_PROTO_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_PROTO_CTOR_FUNC>::type>();
                case cantor::pb::ast::ExprNodeType::ENT_UDF_CTOR_FUNC:
                    return cast_to<ExpressionTypeTrait<cantor::pb::ast::ExprNodeType::ENT_UDF_CTOR_FUNC>::type>();
                default:
                    return turbo::invalid_argument_error("unsupported expression type");
            }
            TURBO_UNREACHABLE();
        }

        turbo::Result<ExprVisitor> cast_to() const {
            return cast_to(expr->type());
        }
    };
} // namespace cantor::unresolved
